package org.liu.knowledge.store;

import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.ollama.OllamaEmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.pgvector.PgVectorEmbeddingStore;
import org.testcontainers.containers.PostgreSQLContainer;
import org.testcontainers.utility.DockerImageName;

import java.util.List;

public class PgVectorEmbeddingStoreExample {

    public static void main(String[] args) {


        EmbeddingModel embeddingModel = OllamaEmbeddingModel.builder()
                .baseUrl("http://localhost:11434/")
                .modelName("zyw0605688/gte-large-zh")
                .build();


        EmbeddingStore<TextSegment> embeddingStore = PgVectorEmbeddingStore.builder()
                //指定主机地址
                .host("192.168.10.100")
                //指定端口
                .port(5432)
                //指定用户名
                .user("root")
                //指定密码
                .password("Password123@postgres")
                //指定数据库名
                .database("root")
                //指定向量数据所在表名
                .table("knowledge_vector")
                //指定向量维度
                .dimension(embeddingModel.dimension())
                .createTable(true)
                .build();

//        TextSegment segment1 = TextSegment.from("I like football.");
//        segment1.metadata().put("id",111);
//        Embedding embedding1 = embeddingModel.embed(segment1).content();
//        embeddingStore.add(embedding1, segment1);
//
//        TextSegment segment2 = TextSegment.from("The weather is good today.");
//        Embedding embedding2 = embeddingModel.embed(segment2).content();
//        embeddingStore.add(embedding2, segment2);


        Embedding queryEmbedding = embeddingModel.embed("What is your favourite sport?").content();


        EmbeddingSearchRequest embeddingSearchRequest = EmbeddingSearchRequest.builder()
                .queryEmbedding(queryEmbedding)
                .maxResults(1)
                .build();

        List<EmbeddingMatch<TextSegment>> relevant = embeddingStore.search(embeddingSearchRequest).matches();


        EmbeddingMatch<TextSegment> embeddingMatch = relevant.get(0);

        System.out.println(embeddingMatch.score()); // 0.8144288608390052
        System.out.println(embeddingMatch.embedded().text()); // I like football.

    }
}